package com.ayou.eggguard.farm.model;
import org.junit.Test;
import weka.classifiers.Classifier;
import weka.classifiers.functions.Logistic;
import weka.core.*;
import weka.core.converters.ConverterUtils.DataSource;

import java.util.ArrayList;

public class DiseasePrediction {

    public static void main(String[] args) throws Exception {
        // 加载数据集
        DataSource source = new DataSource("D:/dataset.arff");
        Instances data = source.getDataSet();
        if (data.classIndex() == -1)
            data.setClassIndex(data.numAttributes() - 1);

        // 创建逻辑回归分类器
        Classifier logistic = new Logistic();

        // 训练模型
        logistic.buildClassifier(data);

        // 保存模型
        SerializationHelper.write("D:/dataset.model", logistic);
        // Load the model from the file
        Classifier loadedModel = (Classifier) SerializationHelper.read("D:/dataset.model");

        // 创建新的实例进行预测
        Instance newInst = new DenseInstance(7);
        newInst.setDataset(data);
        newInst.setValue(0, 43.3);
        newInst.setValue(1, 6);
        newInst.setValue(2,2 );
        newInst.setValue(3,3);
        newInst.setValue(4, 4);
        newInst.setValue(5,4);
        newInst.setValue(6, 1);

        double predictedClass = loadedModel.classifyInstance(newInst);
        // 获取类别概率分布
        double[] probabilities = loadedModel.distributionForInstance(newInst);

        // 输出每个类别的概率
        Attribute classAttribute = data.classAttribute();
        for (int i = 0; i < probabilities.length; i++) {
            String className = classAttribute.value(i);
            System.out.println("Probability of class " + className + ": " + probabilities[i]);
        }

        // 输出预测结果
        String predictedClassName = classAttribute.value((int) predictedClass);
        System.out.println("Predicted class: " + predictedClassName);
    }
}
